import cv2
import torch
from paddleocr import PaddleOCR
from src.utils.logger import log_error, log_info


class OCRProcessor:
    def __init__(self, use_angle_cls=True, lang='ch', use_gpu=True):
        """
        初始化OCR处理器
        :param use_angle_cls: 是否使用角度分类
        :param lang: OCR识别语言
        :param use_gpu: 是否使用GPU
        """
        self.paddle_ocr = PaddleOCR(use_angle_cls=use_angle_cls, lang=lang, use_gpu=use_gpu)
    
    def calc_ocr_result(self, ocr_result, confidence_threshold=0.5):
        """
        对图像进行OCR识别，提取文本信息
        :param image_path: 图像路径
        :param confidence_threshold: 置信度阈值
        :return: 识别到的文本列表
        """
        result = ocr_result
        detected_texts = []

        if result is None or len(result) == 0:
            return detected_texts

        for line in result:
            if line is None:
                continue
            for box, (text, confidence) in line:
                if confidence >= confidence_threshold:
                    # 计算边界框的位置特征
                    x_min = min(p[0] for p in box)
                    y_min = min(p[1] for p in box)
                    x_max = max(p[0] for p in box)
                    y_max = max(p[1] for p in box)
                    width = x_max - x_min
                    height = y_max - y_min
                    area = width * height
                    aspect_ratio = width / height if height > 0 else 0

                    detected_texts.append({
                        "text": text,
                        "confidence": confidence,
                        "box": box,
                        "position": {
                            "x_min": x_min,
                            "y_min": y_min,
                            "x_max": x_max,
                            "y_max": y_max,
                            "center_x": sum(p[0] for p in box) / 4,
                            "center_y": sum(p[1] for p in box) / 4,
                            "width": width,
                            "height": height,
                            "area": area,
                            "aspect_ratio": aspect_ratio
                        }
                    })

        return detected_texts
    
    def recognize_filtered_subtitles(self, image_path):
        """从图像中识别过滤后的字幕"""
        # 从原VideoProcessor移植相关代码
        try:
            # OCR识别处理
            result = self.paddle_ocr.ocr(image_path, cls=True)
            if result == [None]:
                log_info(f"OCR识别结果为空: {image_path}")
                return []
            # 过滤和处理OCR结果
            filtered_results = []
            if result and len(result) > 0:
                for line in result[0]:
                    text = line[1][0]
                    confidence = line[1][1]
                    bbox = line[0]
                    
                    # 提取位置信息
                    x1, y1 = bbox[0]
                    x2, y2 = bbox[2]
                    width = x2 - x1
                    height = y2 - y1
                    center_x = (x1 + x2) / 2
                    center_y = (y1 + y2) / 2
                    
                    position = {
                        "x1": x1, "y1": y1, 
                        "x2": x2, "y2": y2,
                        "width": width, "height": height,
                        "center_x": center_x, "center_y": center_y
                    }
                    
                    filtered_results.append({
                        "text": text,
                        "confidence": confidence,
                        "position": position
                    })
            
            return filtered_results
        except Exception as e:
            log_error(f"OCR识别错误: {e}")
            return []

    def extract_subtitle_mask(self, img_bgr):
        hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV)
        gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
        img_h, img_w = img_bgr.shape[:2]

        # 1. 基于HSV的多色彩字幕检测
        white_mask = cv2.inRange(hsv, (0, 0, 180), (180, 40, 255))  # 白色/淡色
        colored_mask = cv2.inRange(hsv, (0, 40, 180), (180, 255, 255))  # 彩色高亮
        yellow_mask = cv2.inRange(hsv, (20, 100, 100), (40, 255, 255))  # 黄色特别处理

        # 2. 基于边缘的文本检测
        # Canny边缘检测可以帮助识别文本边缘
        edges = cv2.Canny(gray, 100, 200)
        edge_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 1))
        edges = cv2.dilate(edges, edge_kernel, iterations=1)

        # 3. 组合所有掩码
        hsv_mask = cv2.bitwise_or(cv2.bitwise_or(white_mask, colored_mask), yellow_mask)

        # 4. 将HSV掩码与边缘掩码结合
        combined_mask = cv2.bitwise_or(hsv_mask, edges)

        # 5. 限制到底部区域 (通常字幕在下方)
        # 取消注释以限制检测区域
        # combined_mask[0:int(img_h * 0.7), :] = 0

        # 6. 结构增强：使用横向结构元素连接字符
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 3))
        final_mask = cv2.dilate(combined_mask, kernel, iterations=1)

        return final_mask
